
from utils.utils import evaluate
import torch.nn as nn
from tqdm.notebook import tqdm
import torch.optim as optim


train_param = {
    "epoch" : 50,
    'model_lr': 1e-3,
    'criterion' : nn.CrossEntropyLoss(),
    'criterion_MSE' : nn.MSELoss(),

}

def train(model, dataloader, clean_testset_loader, backdoor_testset_loader, param, train_param):
    optimizer = optim.Adam(model.parameters(), lr=train_param['model_lr'])
    criterion = train_param['criterion']
    model.train()
    # 训练模型
    num_epochs = train_param['epoch'] # 增加训练轮数
    for epoch in range(num_epochs):
        for data in tqdm(dataloader, desc="Training"): # for data in mixed_dataloader:
            if len(data) == 3:
                imgs, labels, is_bd = data
            else:
                imgs, labels = data
                
            imgs, labels = imgs.to(param.device), labels.to( param.device)
                
            out = model(imgs)
            loss = criterion(out, labels)
               
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
            
        
        avg_clean_loss, clean_acc = evaluate(model, clean_testset_loader, criterion,  param.device)
        avg_backdoor_loss, backdoor_acc = evaluate(model, backdoor_testset_loader, criterion,  param.device)
        print("Deep:\nclean_acc:{} , backdoor_acc:{}".format(clean_acc,backdoor_acc))
        print("avg_clean_loss:{} , avg_backdoor_loss:{}\n".format(avg_clean_loss,avg_backdoor_loss))
        
def evaluate(model, dataloader, criterion, device):
    model.eval()  # 设置模型为评估模式
    correct = 0
    total = 0
    total_loss = 0.0
    with torch.no_grad():  # 关闭梯度计算
        for data in tqdm(dataloader, desc="Evaluating Model"):
            # Handle the case where data has 2 or 3 elements
            if len(data) == 2:
                imgs, labels = data
            else:
                imgs, labels, _ = data
            imgs = imgs.to(device)
            labels = labels.to(device)
            
            outputs = model(imgs)
            loss = criterion(outputs, labels)
            total_loss += loss.item() * labels.size(0)
            
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    avg_loss = total_loss / total
    accuracy = correct / total
    model.train()  # 重新设置模型为训练模式
    return avg_loss, accuracy